package com.fliad.ai.service;

import cn.hutool.core.bean.BeanUtil;
import cn.hutool.core.util.IdUtil;
import cn.hutool.core.util.StrUtil;
import com.fliad.ai.dto.AiMcpDto;
import com.fliad.ai.entity.AiMcp;
import org.noear.snack.ONode;
import org.noear.solon.ai.mcp.client.McpClientProvider;
import org.noear.solon.ai.mcp.client.McpServerParameters;
import org.noear.solon.annotation.Component;
import org.noear.solon.annotation.Init;
import org.noear.wood.DbContext;
import org.noear.wood.annotation.Db;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.noear.solon.core.util.ResourceUtil;

import java.math.BigInteger;
import java.sql.SQLException;
import java.util.*;

@Component
public class AiMcpService {
    @Db
    DbContext db;

    private static final Map<String, McpClientProvider> mcpClientProviderMap = new HashMap<>();
    private static final ObjectMapper objectMapper = new ObjectMapper();

    // 初始化MCP服务
    public void initializeMcpServices() throws SQLException {
        List<AiMcp> mcpList = db.table("ai_mcp")
                .selectList("*", AiMcp.class);

        for (AiMcp mcp : mcpList) {
            AiMcpDto aiMcpDto = convertToDto(mcp);
            addMcpClientProvider(aiMcpDto);
        }
    }

    // 将AiMcp对象转换为AiMcpDto对象
    private AiMcpDto convertToDto(AiMcp mcp) {
        AiMcpDto aiMcpDto = ONode.load(mcp).toObject(AiMcpDto.class);
        aiMcpDto.setEnvVars(!StrUtil.isEmptyIfStr(aiMcpDto.getEnvVars()) ? ONode.loadStr(mcp.getEnvVars()) : null);
        return aiMcpDto;
    }

    // 添加MCP客户端提供者
    private void addMcpClientProvider(AiMcpDto aiMcpDto) {
        String name = aiMcpDto.getName();
        String type = aiMcpDto.getType();
        String command = aiMcpDto.getCommand();
        String args = aiMcpDto.getArgs();

        McpClientProvider mcpClientProvider;
        if ("stdio".equals(type)) {
            mcpClientProvider = McpClientProvider.builder().channel(type).serverParameters(
                    McpServerParameters.builder(command).args(args.split(" "))
                            .env(aiMcpDto.getEnvVars() != null ? aiMcpDto.getEnvVars().toObject() : Collections.emptyMap())
                            .build()
            ).build();
        } else {
            mcpClientProvider = McpClientProvider.builder().channel(type).apiUrl(aiMcpDto.getServiceUrl()).build();
        }

        mcpClientProviderMap.put(name, mcpClientProvider);
        mcpClientProvider.getClient().initialize();
    }

    // 创建MCP服务
    public boolean createMcp(AiMcp mcp) throws SQLException {
        mcp.setCreated(new Date());
        mcp.setModified(new Date());
        boolean result = db.table("ai_mcp")
                .setEntity(mcp)
                .insert() > 0;
        if (result) {
            addMcpClientProvider(convertToDto(mcp));
        }
        return result;
    }

    // 获取MCP客户端提供者
    public McpClientProvider getMcpClientProvider(String name) {
        return mcpClientProviderMap.get(name);
    }

    // 销毁McpClientProvider
    public void destroyMcpClientProvider(String name) {
        McpClientProvider mcpClientProvider = mcpClientProviderMap.remove(name);
        if (mcpClientProvider != null) {
            mcpClientProvider.close();
        }
    }

    // 删除MCP服务
    public boolean deleteMcp(BigInteger id) throws SQLException {
        AiMcp mcp = db.table("ai_mcp")
                .whereEq("id", id)
                .selectItem("*", AiMcp.class);
        if (mcp != null) {
            destroyMcpClientProvider(mcp.getName());
        }
        return db.table("ai_mcp")
                .whereEq("id", id)
                .delete() > 0;
    }

    // 更新MCP服务
    public boolean updateMcp(Long id, AiMcp mcp) throws SQLException {
        mcp.setModified(new Date());
        boolean result = db.table("ai_mcp")
                .whereEq("id", id)
                .setEntity(mcp)
                .update() > 0;
        if (result) {
            destroyMcpClientProvider(mcp.getName());
            addMcpClientProvider(convertToDto(mcp));
        }
        return result;
    }

    // 获取MCP服务列表
    public List<AiMcpDto> getMcpList() throws SQLException {
        List<AiMcp> mcpList = db.table("ai_mcp")
                .selectList("*", AiMcp.class);

        List<AiMcpDto> mcpDtoList = new ArrayList<>();
        for (AiMcp mcp : mcpList) {
            AiMcpDto mcpDto = convertToDto(mcp);
            updateMcpConnectionStatus(mcpDto);
            mcpDtoList.add(mcpDto);
        }

        return mcpDtoList;
    }

    // 更新MCP连接状态
    private void updateMcpConnectionStatus(AiMcpDto mcpDto) {
        McpClientProvider mcpClientProvider = getMcpClientProvider(mcpDto.getName());
        if (mcpClientProvider != null) {
            try {
                mcpDto.setConnectionStatus(mcpClientProvider.getClient().isInitialized() ? "connected" : "disconnected");
            } catch (Throwable e) {
                mcpDto.setConnectionStatus("disconnected");
            }
            updateMcpTools(mcpDto, mcpClientProvider);
        } else {
            mcpDto.setConnectionStatus("disconnected");
        }
    }

    // 更新MCP工具信息
    private void updateMcpTools(AiMcpDto mcpDto, McpClientProvider mcpClientProvider) {
        ONode tools = ONode.newArray();
        mcpClientProvider.getTools().stream().forEach(functionTool -> {
            tools.addNode(ONode.newObject()
                    .set("toolName", functionTool.name())
                    .set("description", functionTool.description())
                    .set("inputSchema", ONode.loadStr(functionTool.inputSchema()))
            );
        });
        mcpDto.setTools(tools);
    }

    // 调用MCP工具
    public String callMcpTool(String mcpName, String toolName, Map<String, Object> inputParams) {
        McpClientProvider mcpClientProvider = getMcpClientProvider(mcpName);
        if (mcpClientProvider == null) {
            throw new RuntimeException("MCP服务不存在: " + mcpName);
        }

        return mcpClientProvider.callToolAsText(toolName, inputParams).getContent();
    }

    // 查询单个 MCP 服务（从mcp_square.json读取）
    public AiMcp getSquareMcp(String id) throws Exception {
        JsonNode jsonNode = readMcpSquareJson();
        JsonNode serverNode = jsonNode.get("mcpServers").get(id);

        if (serverNode == null) {
            throw new RuntimeException("MCP服务不存在");
        }

        AiMcp mcp = new AiMcp();
        mcp.setName(id);
        mcp.setType(serverNode.has("type") ? serverNode.get("type").asText() : "STDIO");
        mcp.setCommand(serverNode.has("command") ? serverNode.get("command").asText() : "");
        mcp.setArgs(buildArgsString(serverNode.get("args")));
        // 处理环境变量
        if (serverNode.has("env")) {
            mcp.setEnvVars(buildEnvVarsString(serverNode.get("env")));
        }

        return mcp;
    }

    // 读取mcp_square.json文件
    private JsonNode readMcpSquareJson() throws Exception {
        return objectMapper.readTree(ResourceUtil.findResourceOrFile(null, "./static/mcp_square.json"));
    }

    // 构建参数字符串
    private String buildArgsString(JsonNode argsNode) {
        if (argsNode == null || !argsNode.isArray()) {
            return "";
        }
        StringBuilder argsBuilder = new StringBuilder();
        for (JsonNode argNode : argsNode) {
            if (argsBuilder.length() > 0) {
                argsBuilder.append(" ");
            }
            argsBuilder.append(argNode.asText());
        }
        return argsBuilder.toString();
    }

    // 构建环境变量字符串
    private String buildEnvVarsString(JsonNode envNode) {
        if (envNode == null || !envNode.isObject()) {
            return "";
        }
        StringBuilder envVarsBuilder = new StringBuilder();
        Iterator<Map.Entry<String, JsonNode>> envIterator = envNode.fields();
        while (envIterator.hasNext()) {
            Map.Entry<String, JsonNode> envEntry = envIterator.next();
            if (envVarsBuilder.length() > 0) {
                envVarsBuilder.append(" ");
            }
            envVarsBuilder.append(envEntry.getKey()).append("=").append(envEntry.getValue().asText());
        }
        return envVarsBuilder.toString();
    }

    // 查询 MCP 服务列表（从mcp_square.json读取）
    public List<AiMcpDto> listSquareMcp() throws Exception {
        JsonNode jsonNode = readMcpSquareJson();
        JsonNode serversNode = jsonNode.get("mcpServers");

        if (serversNode == null || !serversNode.isObject()) {
            throw new RuntimeException("无效的MCP配置文件");
        }

        List<AiMcpDto> mcpList = new ArrayList<>();
        Iterator<Map.Entry<String, JsonNode>> iterator = serversNode.fields();

        while (iterator.hasNext()) {
            Map.Entry<String, JsonNode> entry = iterator.next();
            String id = entry.getKey();
            JsonNode serverNode = entry.getValue();

            AiMcpDto mcp = new AiMcpDto();
            mcp.setName(id);
            mcp.setType(serverNode.has("command") ? "stdio" : "sse");
            mcp.setServiceUrl(serverNode.has("url") ? serverNode.get("url").asText() : "");
            mcp.setCommand(serverNode.has("command") ? serverNode.get("command").asText() : "");
            mcp.setArgs(buildArgsString(serverNode.get("args")));
            mcp.setEnvVars(buildEnvVarsObject(serverNode.get("env")));

            mcpList.add(mcp);
        }

        return mcpList;
    }

    // 构建环境变量对象
    private ONode buildEnvVarsObject(JsonNode envNode) {
        if (envNode == null || !envNode.isObject()) {
            return null;
        }
        ONode envVars = ONode.newObject();
        Iterator<Map.Entry<String, JsonNode>> envIterator = envNode.fields();
        while (envIterator.hasNext()) {
            Map.Entry<String, JsonNode> envEntry = envIterator.next();
            envVars.set(envEntry.getKey(), envEntry.getValue().asText());
        }
        return envVars;
    }

    // 新增方法：校验安装参数
    private void validateInstallParams(Map<String, Object> params) throws SQLException {
        String name = (String) params.get("name");
        if (StrUtil.isEmpty(name)) {
            throw new IllegalArgumentException("服务名称不能为空");
        }

        // 校验名称是否已存在
        AiMcp existingMcp = db.table("ai_mcp")
                .whereEq("name", name)
                .selectItem("*", AiMcp.class);
        if (existingMcp != null && existingMcp.getId() != null) {
            throw new RuntimeException("服务名称已存在: " + name);
        }

        String type = (String) params.get("type");
        if (!"stdio".equals(type) && !"sse".equals(type)) {
            throw new IllegalArgumentException("不支持的服务类型: " + type);
        }
    }

    // 新增方法：创建 McpClientProvider
    private McpClientProvider createMcpClientProvider(AiMcpDto mcp) {
        String type = mcp.getType();
        if ("stdio".equals(type)) {
            return McpClientProvider.builder()
                    .channel(type)
                    .serverParameters(McpServerParameters.builder(mcp.getCommand())
                            .args(mcp.getArgs().split(" "))
                            .env(mcp.getEnvVars() != null ? mcp.getEnvVars().toObject() : Collections.emptyMap())
                            .build())
                    .build();
        } else {
            return McpClientProvider.builder()
                    .channel(type)
                    .apiUrl(mcp.getServiceUrl())
                    .build();
        }
    }

    // 新增方法：将 MCP 信息保存到数据库
    private boolean saveMcpToDatabase(AiMcpDto mcpDto) throws SQLException {
        AiMcp aiMcp = BeanUtil.toBean(mcpDto, AiMcp.class);
        return db.table("ai_mcp")
                .setEntity(aiMcp)
                .insert() > 0;
    }

    // 优化后的 install 方法
    public boolean install(Map<String, Object> params) throws SQLException {
        // 校验参数
        validateInstallParams(params);

        // 构建 MCP 对象
        AiMcpDto mcp = new AiMcpDto();
        mcp.setId(IdUtil.getSnowflakeNextId());
        mcp.setName((String) params.get("name"));
        mcp.setType((String) params.get("type"));
        mcp.setCommand((String) params.get("command"));
        mcp.setArgs((String) params.get("args"));
        mcp.setServiceUrl((String) params.get("serviceUrl"));
        if (params.get("envVars") != null) {
            mcp.setEnvVars(ONode.load(params.get("envVars")));
        }
        mcp.setCreated(new Date());
        mcp.setModified(new Date());

        // 使用 addMcpClientProvider 方法创建并初始化 McpClientProvider
        addMcpClientProvider(mcp);

        // 保存到数据库
        return saveMcpToDatabase(mcp);
    }

    // 新增方法：根据ID查询单个MCP服务
    public AiMcpDto getMcpById(BigInteger id) throws SQLException {
        AiMcp mcp = db.table("ai_mcp")
                .whereEq("id", id)
                .selectItem("*", AiMcp.class);
        if (mcp == null) {
            throw new RuntimeException("MCP服务不存在");
        }
        AiMcpDto mcpDto = convertToDto(mcp);
        updateMcpConnectionStatus(mcpDto);
        return mcpDto;
    }

}
